#pragma once

#include <vector>
#include <cassert>
#include <ostream>
#include <stdint.h>
#include <Eigen/Dense>

namespace  ldl
{
namespace core
{
template<typename T>
class Tensor;

template<typename T>
class TensorView
{
public:
    TensorView(const std::vector<int64_t> &shape, T *data)
    {
        m_data = data;
        m_shape = shape;
        m_size = 1;
        for(auto dim : shape) {
            m_size *= dim;
        }
    }

    TensorView(const Tensor<T> &tensor)
    {
        m_data = (T *)tensor.data().data();
        m_shape = tensor.shape();
    }

    TensorView<T> operator[](int64_t index) const
    {
        auto stride = m_size / m_shape.at(0);
        std::vector<int64_t> shape(m_shape.begin() + 1, m_shape.begin() + m_shape.size());
        if(m_shape.size() == 1) {
            shape.push_back(1);
        }
        return TensorView<T>(shape, m_data + index * stride);
    }

    static void dot(const TensorView &a, const TensorView &b, TensorView &out)
    {
        assert(a.m_shape.size() == 2);
        assert(b.m_shape.size() == 2);
        assert(a.m_shape.at(1) ==  b.m_shape.at(0));
        assert(out.m_shape.size() == 2);
        assert(out.m_shape[0] == a.m_shape[0]);
        assert(out.m_shape[1] == b.m_shape[1]);

        Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>> mat_a(a.m_data, a.m_shape.at(0), a.m_shape.at(1));
        Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>> mat_b(b.m_data, b.m_shape.at(0), b.m_shape.at(1));
        Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>> mat_out(out.m_data, out.m_shape.at(0), out.m_shape.at(1));

        mat_out.noalias() = mat_a * mat_b;
    }

    void dump(std::ostream& os, const TensorView<T>& tensor_view) const
    {
        os << "[";
        if(tensor_view.m_shape.size() == 1)
        {

            auto dim  = tensor_view.m_shape.at(0);
            for(int64_t index = 0; index < dim; index++)
            {
                os << tensor_view.m_data[index];
                if(index != (dim - 1))
                    os << ",";
            }
        }
        else
        {
            auto dim  = tensor_view.m_shape.at(0);
            for(int64_t index = 0; index < dim; index++)
            {
                dump(os, tensor_view[index]);
                if(index != (dim - 1))
                {
                    os << ",\n";
                }

                if((tensor_view.m_shape.size() > 2) && (dim != index + 1))
                {
                    os << "\n";
                }
            }
        }
        os << "]";
    }

    friend std::ostream& operator<<(std::ostream& os, const TensorView<T>& tensor_view)
    {
        os << "Tensor(\n";
        tensor_view.dump(os, tensor_view);
        os << ",shape=(";
        for (size_t i = 0; i < tensor_view.m_shape.size(); ++i) {
            if (i != 0) os << ",";
            os << tensor_view.m_shape[i];
        }
        os << "))";
        return os;
    }

    std::vector<int64_t> shape()
    {
        return m_shape;
    }

    int64_t size()
    {
        return m_size;
    }
private:
    T *m_data;
    int64_t m_size;
    std::vector<int64_t> m_shape;
};
}
}
